import configparser

from eventlet import Semaphore

from managers.nodes.baseGraphNode import  *
from agent.common_base_agent import CommonBaseAgent
from ..execute.baseExeHis import BaseExeHis
from .sessionInfo import SessionInfo
from collections import defaultdict

class AgentGraphNode(BaseGraphNode):
    def __init__(self, id:str, data:dict, graphNodeManager, random_key:str):
        super().__init__(id, data, graphNodeManager, random_key)
        self.name = data['data']['name']
        self.role = data['data'].get('role', 'assistant')
        self.prompt = data['data'].get('prompt', 'You are a helpful assistant.')
        self.isSpeak = data['data'].get('isSpeak', True)
        self.isFunctionCall = data['data'].get('isFunctionCall', False)
        self.customKey = None
        if data['data'].get('has_custom_key', False):
            self.customKey = data['data']['custom_key'] # {'key':{'name':'name', 'type':'type', 'description':'description'}}
        # self.agent = CommonBaseAgent(name = self.name, sys_prompt = self.prompt, model_config_name = 'openai_cfg', isSpeak = self.isSpeak, isFunctionCall = self.isFunctionCall, customKey = self.customKey)
        self.session_locks = defaultdict(Semaphore)  # 每个session_id独立锁
        self.uuid2agent = {}


    def execute(self, session_info: SessionInfo=None, exe_his: BaseExeHis=None, trigger=None, running_id=None):
        if session_info is None: return 
        if session_info.id not in self.uuid2agent:
            self.uuid2agent[session_info.id] = CommonBaseAgent(name = self.role, sys_prompt = self.prompt, model_config_name = 'openai_cfg', isSpeak = self.isSpeak, isFunctionCall = self.isFunctionCall, customKey = self.customKey)
        agent = self.uuid2agent[session_info.id]
        query = ""
        if 'query-msg' in self.dict['in']: # {'nodeId': sourceID, 'handle': sourceHandle}
            queryFrom = self.dict['in']['query-msg']
            prenode =self.graphNodeManager.nodesByIdDict[queryFrom['nodeId']]
            query = prenode.get_newest_output_in_session(session_info, queryFrom['handle'])
        
        if 'context-msg' in self.dict['in']:
            contextFrom = self.dict['in']['context-msg']
            prenode = self.graphNodeManager.nodesByIdDict[contextFrom['nodeId']]
            context = prenode.get_newest_output_in_session(session_info, contextFrom['handle'])
            self.agent.context = context

        if 'memory-msgs' in self.dict['in']:
            memoryFrom = self.dict['in']['memory-msgs']
            prenode = self.graphNodeManager.nodesByIdDict[memoryFrom['nodeId']]
            memorys = prenode.get_newest_output_in_session(session_info, memoryFrom['handle'])
            self.agent.setmemory(memorys)
        session_info.ws.start_background_task(self._process_agent_reply, session_info, query, exe_his)

        

    def _process_agent_reply(self, session_info: SessionInfo, query, exe_his, trigger, running_id):
        returnMsg = self.agent.reply({"role": "user", "content": query}, exe_his)
        with session_info.out_write_lock:
            if session_info.out.get(self.id,None) is None: session_info.out[self.id] = {}
            returnMsg = returnMsg['content']
            if('reason' in returnMsg):
                if 'reason-msg' not in session_info.out[self.id]: session_info.out[self.id]['reason-msg'] = []
                session_info.out[self.id]['reason-msg'].append(returnMsg['reason'])

        self.executed_num += 1
        if exe_his:exe_his.done()
        if trigger and trigger.enInterrupt and trigger.session_id_2_running_id[session_info.id] != running_id: return 
        self.flow_next(session_info, trigger=trigger, running_id=running_id, type='text', role='assistant')
    def flow_next(self, session_info: SessionInfo, **kwargs):
        return super().flow_next(session_info, **kwargs)